"""
MCP客户端
集成所有MCP功能到单一文件中，包括：
- MCP服务器连接和健康检测
- 工具获取和调用
- 后台重连机制
"""

import asyncio
import json
import logging
import socket
import random
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional, Callable
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import urlparse
from contextlib import AsyncExitStack

# MCP相关导入
try:
    import mcp
    from mcp.client.sse import sse_client
    from mcp.client.streamable_http import streamablehttp_client
    MCP_AVAILABLE = True
except ImportError as e:
    logging.warning(f"MCP库不可用: {e}")
    MCP_AVAILABLE = False

# 安全导入Function Calling模块
try:
    from utils.function_call.function_registry import FunctionRegistry
    from utils.function_call.function_schema import get_all_schemas
    LOCAL_TOOLS_AVAILABLE = True
except ImportError:
    LOCAL_TOOLS_AVAILABLE = False
    get_all_schemas = lambda: []

logger = logging.getLogger(__name__)

@dataclass
class MCPTool:
    """MCP工具描述"""
    name: str
    parameters: Dict[str, Any]
    description: str
    server_name: str
    active: bool = True

@dataclass
class MCPServerState:
    """MCP服务器状态"""
    name: str
    config: dict
    status: str = "disconnected"  # disconnected, connecting, connected, failed
    last_health_check: Optional[datetime] = None
    retry_count: int = 0
    next_retry_time: Optional[datetime] = None
    tools: List[MCPTool] = None
    error_message: Optional[str] = None

    def __post_init__(self):
        if self.tools is None:
            self.tools = []

    def should_retry_now(self) -> bool:
        """检查是否应该现在重试连接"""
        if self.status == "connected":
            return False
        if self.next_retry_time is None:
            return True
        return datetime.now() >= self.next_retry_time

class MCPConnection:
    """单个MCP连接"""

    def __init__(self, name: str, config: dict):
        self.name = name
        self.config = config
        self.session: Optional[mcp.ClientSession] = None
        self.exit_stack = AsyncExitStack()
        self.tools: List[mcp.Tool] = []
        self._streams_context = None

    async def health_check(self, timeout: float = 3.0) -> bool:
        """快速健康检查"""
        if "url" not in self.config:
            return True  # stdio连接无法网络检查

        try:
            parsed_url = urlparse(self.config["url"])
            host = parsed_url.hostname
            port = parsed_url.port or (443 if parsed_url.scheme == "https" else 80)

            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            result = sock.connect_ex((host, port))
            sock.close()

            return result == 0
        except Exception:
            return False

    async def connect(self) -> None:
        """连接到MCP服务器"""
        if not MCP_AVAILABLE:
            raise RuntimeError("MCP库不可用")

        cfg = self.config.copy()
        if "mcpServers" in cfg and len(cfg["mcpServers"]) > 0:
            key_0 = list(cfg["mcpServers"].keys())[0]
            cfg = cfg["mcpServers"][key_0]
        cfg.pop("active", None)

        # 健康检查
        if "url" in cfg and not await self.health_check():
            raise ConnectionError(f"服务器 {self.name} 健康检查失败")

        # 连接
        if "url" in cfg:
            await self._connect_via_url(cfg)
        else:
            await self._connect_via_stdio(cfg)

        await self.session.initialize()
        await self._load_tools()

    async def _connect_via_url(self, cfg: dict) -> None:
        """通过URL连接"""
        transport_type = cfg.get("transport", cfg.get("type", "sse"))
        if transport_type in ["streamable-http", "streamable_http"]:
            timeout = timedelta(seconds=cfg.get("timeout", 30))
            self._streams_context = streamablehttp_client(
                url=cfg["url"],
                headers=cfg.get("headers", {}),
                timeout=timeout,
                sse_read_timeout=timedelta(seconds=cfg.get("sse_read_timeout", 300)),
                terminate_on_close=cfg.get("terminate_on_close", True),
            )
            read_s, write_s, _ = await self._streams_context.__aenter__()
            self.session = await self.exit_stack.enter_async_context(
                mcp.ClientSession(read_stream=read_s, write_stream=write_s)
            )
        else:
            # SSE连接
            self._streams_context = sse_client(
                url=cfg["url"],
                headers=cfg.get("headers", {}),
                timeout=cfg.get("timeout", 5),
                sse_read_timeout=cfg.get("sse_read_timeout", 300),
            )
            streams = await self._streams_context.__aenter__()
            self.session = await self.exit_stack.enter_async_context(
                mcp.ClientSession(*streams)
            )

    async def _connect_via_stdio(self, cfg: dict) -> None:
        """通过stdio连接"""
        server_params = mcp.StdioServerParameters(**cfg)
        stdio_transport = await self.exit_stack.enter_async_context(
            mcp.stdio_client(server_params, errlog=lambda msg: logger.error(f"MCP {self.name}: {msg}"))
        )
        self.session = await self.exit_stack.enter_async_context(
            mcp.ClientSession(*stdio_transport)
        )

    async def _load_tools(self) -> None:
        """加载工具列表"""
        if not self.session:
            return
        response = await self.session.list_tools()
        self.tools = response.tools

    async def call_tool(self, tool_name: str, arguments: dict) -> Any:
        """调用工具"""
        if not self.session:
            raise RuntimeError("连接未建立")
        return await self.session.call_tool(tool_name, arguments)

    async def cleanup(self) -> None:
        """清理连接"""
        try:
            if self._streams_context:
                await self._streams_context.__aexit__(None, None, None)
            await self.exit_stack.aclose()
        except Exception:
            pass

class MCPClient:
    """MCP客户端"""

    def __init__(self):
        self.servers: Dict[str, MCPServerState] = {}
        self.connections: Dict[str, MCPConnection] = {}
        self.background_task: Optional[asyncio.Task] = None
        self.shutdown_event = asyncio.Event()
        self.tools_update_callback: Optional[Callable] = None

    async def initialize(self) -> None:
        """初始化客户端"""
        logger.info("初始化MCP客户端...")

        # 加载配置
        await self._load_config()

        # 启动后台监控
        self.background_task = asyncio.create_task(self._background_monitor())

        logger.info("MCP客户端初始化完成")

    async def _load_config(self) -> None:
        """加载配置文件"""
        config_file = Path(__file__).parent.parent / "config" / "mcp_servers.json"

        if not config_file.exists():
            # 创建默认配置
            config_file.parent.mkdir(parents=True, exist_ok=True)
            default_config = {"mcpServers": {}}
            with open(config_file, "w", encoding="utf-8") as f:
                json.dump(default_config, f, ensure_ascii=False, indent=4)
            return

        try:
            with open(config_file, "r", encoding="utf-8") as f:
                config = json.load(f)

            for name, server_config in config.get("mcpServers", {}).items():
                if server_config.get("active", True):
                    self.servers[name] = MCPServerState(name=name, config=server_config)

        except Exception as e:
            logger.error(f"加载配置失败: {e}")

    async def _background_monitor(self) -> None:
        """后台监控任务"""
        while not self.shutdown_event.is_set():
            try:
                # 尝试连接或重连服务器
                for name, server in self.servers.items():
                    if server.should_retry_now():
                        await self._attempt_connection(name, server)

                # 健康检查已连接的服务器
                await self._health_check_all()

                # 等待30秒
                await asyncio.wait_for(self.shutdown_event.wait(), timeout=30.0)

            except asyncio.TimeoutError:
                continue
            except Exception as e:
                logger.error(f"后台监控异常: {e}")
                await asyncio.sleep(10)

    async def _attempt_connection(self, name: str, server: MCPServerState) -> None:
        """尝试连接服务器"""
        try:
            server.status = "connecting"
            server.retry_count += 1

            logger.info(f"尝试连接MCP服务器: {name} (第{server.retry_count}次)")

            # 清理旧连接
            if name in self.connections:
                await self.connections[name].cleanup()
                del self.connections[name]

            # 建立新连接
            conn = MCPConnection(name, server.config)
            await conn.connect()

            # 连接成功
            self.connections[name] = conn
            server.status = "connected"
            server.last_health_check = datetime.now()
            server.retry_count = 0
            server.next_retry_time = None
            server.error_message = None

            # 更新工具列表
            server.tools = [
                MCPTool(
                    name=tool.name,
                    parameters=tool.inputSchema,
                    description=tool.description,
                    server_name=name
                )
                for tool in conn.tools
            ]

            logger.info(f"MCP服务器 {name} 连接成功，工具: {[t.name for t in server.tools]}")

            # 通知工具更新
            if self.tools_update_callback:
                if asyncio.iscoroutinefunction(self.tools_update_callback):
                    await self.tools_update_callback()
                else:
                    self.tools_update_callback()

        except Exception as e:
            # 连接失败
            server.status = "failed"
            server.error_message = str(e)

            # 计算下次重试时间（指数退避）
            delay = min(300, 10 * (2 ** min(server.retry_count - 1, 5)))
            delay *= random.uniform(0.5, 1.5)  # 添加抖动
            server.next_retry_time = datetime.now() + timedelta(seconds=delay)

            logger.warning(f"MCP服务器 {name} 连接失败: {e}, 下次重试: {server.next_retry_time.strftime('%H:%M:%S')}")

    async def _health_check_all(self) -> None:
        """对所有连接的服务器进行健康检查"""
        for name, server in self.servers.items():
            if server.status == "connected" and name in self.connections:
                try:
                    conn = self.connections[name]
                    if not await conn.health_check(timeout=5.0):
                        raise ConnectionError("健康检查失败")
                    server.last_health_check = datetime.now()
                except Exception as e:
                    logger.warning(f"服务器 {name} 健康检查失败: {e}")
                    server.status = "failed"
                    server.retry_count = 0
                    server.next_retry_time = datetime.now() + timedelta(seconds=30)

    def get_all_tools_schema(self) -> List[Dict[str, Any]]:
        """获取所有工具的schema（本地+MCP）"""
        # 获取本地工具
        local_schemas = []
        if LOCAL_TOOLS_AVAILABLE:
            try:
                local_schemas = get_all_schemas()
            except Exception as e:
                logger.warning(f"无法获取本地工具: {e}")

        # 获取MCP工具
        mcp_schemas = []
        for server in self.servers.values():
            if server.status == "connected":
                for tool in server.tools:
                    if tool.active:
                        schema = {
                            "type": "function",
                            "function": {
                                "name": tool.name,
                                "description": tool.description,
                                "parameters": tool.parameters
                            }
                        }
                        mcp_schemas.append(schema)

        logger.info(f"获取工具schema: 本地 {len(local_schemas)} 个, MCP {len(mcp_schemas)} 个")
        return local_schemas + mcp_schemas

    def get_mcp_tools_only(self) -> List[Dict[str, Any]]:
        """只获取MCP工具的schema"""
        mcp_schemas = []
        for server_name, server in self.servers.items():
            if server.status == "connected":
                for tool in server.tools:
                    if tool.active:
                        schema = {
                            "type": "function",
                            "function": {
                                "name": tool.name,
                                "description": tool.description,
                                "parameters": tool.parameters
                            },
                            "server_name": server_name
                        }
                        mcp_schemas.append(schema)

        logger.info(f"获取MCP工具schema: {len(mcp_schemas)} 个")
        return mcp_schemas

    async def call_tool(self, tool_name: str, **arguments) -> Any:
        """调用工具（自动判断本地或MCP）"""
        # 先检查是否为MCP工具
        for server in self.servers.values():
            if server.status == "connected":
                for tool in server.tools:
                    if tool.name == tool_name and tool.active:
                        conn = self.connections.get(server.name)
                        if conn:
                            return await conn.call_tool(tool_name, arguments)

        # 尝试本地工具
        if LOCAL_TOOLS_AVAILABLE:
            try:
                return await FunctionRegistry.call(tool_name, **arguments)
            except Exception as e:
                logger.warning(f"本地工具调用失败: {e}")

        raise ValueError(f"工具 '{tool_name}' 不存在或不可用")

    def get_server_status(self) -> Dict[str, Any]:
        """获取服务器状态摘要"""
        connected_count = sum(1 for s in self.servers.values() if s.status == "connected")
        total_tools = sum(len(s.tools) for s in self.servers.values() if s.status == "connected")

        servers_info = []
        for server in self.servers.values():
            info = {
                "name": server.name,
                "status": server.status,
                "tools_count": len(server.tools),
                "last_health_check": server.last_health_check.isoformat() if server.last_health_check else None,
                "error_message": server.error_message
            }
            if server.next_retry_time:
                info["next_retry_time"] = server.next_retry_time.isoformat()
            servers_info.append(info)

        return {
            "mcp_available": MCP_AVAILABLE,
            "total_servers": len(self.servers),
            "connected_servers": connected_count,
            "total_tools": total_tools,
            "servers": servers_info
        }

    async def cleanup(self) -> None:
        """清理资源"""
        logger.info("清理MCP客户端...")

        # 停止后台任务
        self.shutdown_event.set()
        if self.background_task:
            try:
                await asyncio.wait_for(self.background_task, timeout=5.0)
            except asyncio.TimeoutError:
                self.background_task.cancel()

        # 清理所有连接
        for conn in self.connections.values():
            await conn.cleanup()

        self.connections.clear()
        self.servers.clear()

        logger.info("MCP客户端清理完成")

# 全局实例
_mcp_client: Optional[MCPClient] = None

async def get_mcp_client() -> MCPClient:
    """获取全局MCP客户端实例"""
    global _mcp_client
    if _mcp_client is None:
        _mcp_client = MCPClient()
        await _mcp_client.initialize()
    return _mcp_client

async def get_enhanced_function_schemas() -> List[Dict[str, Any]]:
    """获取增强的函数schema列表"""
    client = await get_mcp_client()
    return client.get_all_tools_schema()

async def get_mcp_tools_schemas() -> List[Dict[str, Any]]:
    """获取仅MCP工具的schema列表"""
    client = await get_mcp_client()
    return client.get_mcp_tools_only()

async def enhanced_function_call(tool_name: str, **kwargs) -> Any:
    """增强的函数调用"""
    client = await get_mcp_client()
    return await client.call_tool(tool_name, **kwargs)

async def cleanup_mcp_client() -> None:
    """清理全局MCP客户端"""
    global _mcp_client
    if _mcp_client:
        await _mcp_client.cleanup()
        _mcp_client = None